__author__ = 'KErmakov'
###############################################
##     OVAL Content management library       ##
##          By Kirill Ermakov                ##
##                                           ##
##              OVAL FORGE                   ##
###############################################

# Table of imports
import pymongo
import collections
import functools
from heapq import nsmallest
from operator import itemgetter
import re
import xml.etree.cElementTree as etree
import datetime
import os
from oval_player import oval_mapper
import hashlib
import time
import urllib2
import codecs
## Encoding set
import sys
reload(sys)
sys.setdefaultencoding('utf-8')
# Database connectors
# MongoDB
class mongo_connector():
    def __init__(self, host, port):
        self.db = None
        self.connection = None
        if host and isinstance(host, basestring) and port and isinstance(port, int):
            self.connection = pymongo.Connection(host, port)
            self.db = self.connection['oval-forge']
            self.db.authenticate('oval', 'P@ssw0rd')
# Support class
# Counter class
class Counter(dict):
    def __missing__(self, key):
        return 0
# Hashable dict
class FrozenDict(collections.Mapping):
    """Don't forget the docstrings!!"""

    def __init__(self, *args, **kwargs):
        self._d = dict(*args, **kwargs)

    def __iter__(self):
        return iter(self._d)

    def __len__(self):
        return len(self._d)

    def __getitem__(self, key):
        return self._d[key]

    def __hash__(self):
        return hash(tuple(sorted(self._d.iteritems())))
# Function time calculation decorator for testing
def functime(fn):
    def wrapper(*args, **kwds):
        starttime = time.time()
        result = fn(*args, **kwds)
        ttl = time.time() - starttime
        print(fn.__name__ + " " +  str(ttl))
        return result
    return wrapper
# LFU Caching function
def lfu_cache(maxsize=100):

    def decorating_function(user_function):
        cache = {}                      # mapping of args to results
        use_count = Counter()           # times each key has been accessed
        kwd_mark = object()             # separate positional and keyword args

        @functools.wraps(user_function)
        def wrapper(*args, **kwds):
            key = args
            if kwds:
                key += (kwd_mark,) + tuple(sorted(kwds.items()))
            use_count[key] += 1

            # get cache entry or compute if not found
            try:
                result = cache[key]
                wrapper.hits += 1
            except KeyError:
                result = user_function(*args, **kwds)
                cache[key] = result
                wrapper.misses += 1

                # purge least frequently used cache entry
                if len(cache) > maxsize:
                    for key, _ in nsmallest(maxsize // 10,
                                            use_count.iteritems(),
                                            key=itemgetter(1)):
                        del cache[key], use_count[key]

            return result

        def clear():
            cache.clear()
            use_count.clear()
            wrapper.hits = wrapper.misses = 0

        wrapper.hits = wrapper.misses = 0
        wrapper.clear = clear
        return wrapper
    return decorating_function
# Main classes for OVAL objects
class variable(object):
    def __init__(self, tag=str(), body = str(), attributes = None, variables = None):
        if not attributes:
            self.attributes = dict()
        if not variables:
            self.variables = list()
        if attributes and not isinstance(attributes, dict):
            attributes = dict(attributes)
        if variables and not isinstance(variables, list):
            variables = list(variables)
        self.tag = tag
        self.body = body
        if attributes:
            self.attributes = attributes
        if variables:
            self.variables = variables
        self.hash = 0
        
    def clearme(self):
        hashmap = dict()
        varlist = list(self.variables)
        for var in varlist:
            var.clearme()
            var_hash = var.hashme()
            if var_hash not in hashmap.keys():
                hashmap[var_hash] = None
            else:
                self.variables.remove(var)
    
    def hashme(self):
        summ_hash = hashlib.md5()
        
        for attribute in self.__dict__.keys():
            if attribute != 'variables' and attribute != 'hash' and attribute != 'attributes':
                summ_hash.update(str(self.__getattribute__(attribute)))
        
        if self.attributes:
            for attribute in self.attributes.keys():
                summ_hash.update(str(attribute))
                summ_hash.update(str(self.attributes[attribute]))

        if self.variables:
            for var in self.variables:
                if isinstance(var, variable):
                    if var != self:
                        summ_hash.update(var.hashme())
                else:
                    raise Exception("Error input parameters in variables. Self append ?")
                
        self.hash = summ_hash.hexdigest()
        return self.hash
    
    def serializeme(self):
        self.clearme()
        serialized_object = dict()
        for attribute in self.__dict__.keys():
            if attribute != 'variables:' and attribute != 'hash':
                serialized_object[attribute] = self.__getattribute__(attribute)
        serialized_object['variables'] = list()
        for var in self.variables:
            serialized_object['variables'].append(var.serializeme())
        return serialized_object
# UNIVERSAL OVAL OBJECTS
class oval_object(object):
    def __init__(self):
        self.id = str()
        self.tag = "object"
        self.namespace = str()
        self.oval = str()
        self.version = 0
        self.hash = 0
        self.deprecated = False
        self.variables = list()
        self.notes = list()
        self.comment = str()
        self.signature = str()

    def assign_id(self, UniqId):
        self.id = "oval:"+self.oval+":obj:"+str(UniqId)
        return self.id

    def hashme(self):
        summ_hash = hashlib.md5()

        for attribute in self.__dict__.keys():
            if attribute != 'variables' and attribute != 'hash' and attribute != 'attributes' and attribute!= 'notes' and attribute != 'id' and attribute !='version':
                summ_hash.update(str(self.__getattribute__(attribute)))
        
        for var in self.variables:
            if isinstance(var, variable):
                summ_hash.update(var.hashme())
            else:
                raise Exception("Error input parameters in object variables")
            
        return summ_hash.hexdigest()

    def clearme(self):
        hashmap = dict()
        varlist = list(self.variables)
        for var in varlist:
            var.clearme()
            var_hash = var.hashme()
            if var_hash not in hashmap.keys():
                hashmap[var_hash] = None
            else:
                self.variables.remove(var)
                
    def serializeme(self):
        self.clearme()
        serialized_object = dict()
        for attribute in self.__dict__.keys():
            if attribute != 'variables:' and attribute != 'hash':
                serialized_object[attribute] = self.__getattribute__(attribute)
        serialized_object['variables'] = list()
        for var in self.variables:
            serialized_object['variables'].append(var.serializeme())
        return serialized_object
class oval_state(oval_object):
    def __init__(self):
        super(oval_state, self).__init__()
        self.tag = "state"
        self.operator = str()
        
    def assign_id(self, UniqId):
        self.id = "oval:"+self.oval+":ste:"+str(UniqId)
        return self.id
class oval_variable(oval_object):
    def __init__(self):
        super(oval_variable, self).__init__()
        self.tag = "variable"
        self.datatype = str()

    def assign_id(self, UniqId):
        self.id = "oval:"+self.oval+":var:"+str(UniqId)
        return self.id
class oval_test(oval_object):
    def __init__(self):
        super(oval_test, self).__init__()
        self.tag = "test"
        self.check_existance = str()
        self.check = str()
        self.state_operator = str()
        
    def assign_id(self, UniqId):
        self.id = "oval:"+self.oval+":tst:"+str(UniqId)
        return self.id
class oval_definition(oval_object):
    
    def __init__(self):
        super(oval_definition, self).__init__()
        self.tag = "definition"
        self.defclass = ""
        self.title = ""
        self.description = ""
        self.metadata = list()
        self.criteria = list()
        
    def assign_id(self, UniqId):
        self.id = "oval:"+self.oval+":def:"+str(UniqId)
        return self.id

    def hashme(self):
        summ_hash = hashlib.md5()
        # Dognail for situation, where titles are the same...
        if hasattr(self, 'hiddentitle') and self.hiddentitle and isinstance(self.hiddentitle, basestring):
            summ_hash.update(str(self.hiddentitle))
        else:
            summ_hash.update(str(self.title))
        self.hash = summ_hash.hexdigest()
        return self.hash

    def datahash(self):
        summ_hash = hashlib.md5()
        summ_hash.update(str(self.description))
        for element in self.metadata:
            summ_hash.update(element.hashme())
        for element in self.criteria:
            summ_hash.update(element.hashme())
        return summ_hash.hexdigest()

    def clearme(self):

        hashmap = dict()
        varlist = self.variables
        for var in varlist:
            var.clearme()
            var_hash = var.hashme()
            if var_hash not in hashmap.keys():
                hashmap[var_hash] = None
            else:
                self.variables.remove(var)

        hashmap = dict()
        varlist = list(self.metadata)
        for var in varlist:
            var.clearme()
            var_hash = var.hashme()
            if var_hash not in hashmap.keys():
                hashmap[var_hash] = None
            else:
                self.metadata.remove(var)

        hashmap = dict()
        varlist = list(self.criteria)
        for var in varlist:
            var.clearme()
            var_hash = var.hashme()
            if var_hash not in hashmap.keys():
                hashmap[var_hash] = None
            else:
                self.criteria.remove(var)

    def serializeme(self):
        self.clearme()
        serialized_object = dict()
        for attribute in self.__dict__.keys():
            if attribute != 'variables:' and attribute != 'hash' and attribute != 'metadata' and attribute != 'criteria':
                serialized_object[attribute] = self.__getattribute__(attribute)
                
        serialized_object['metadata'] = list()
        for var in self.metadata:
            serialized_object['metadata'].append(var.serializeme())
        serialized_object['criteria'] = list()
        for var in self.criteria:
            serialized_object['criteria'].append(var.serializeme())

        return serialized_object
# OVAL Forge ------------------------------------------->>>>>>> ENGINE
class forge():

    def __init__(self, host, port, namespace = None):
        self.connector = mongo_connector(host, port)
        self.db = self.connector.db
        if not namespace:
            self.namespace = 'com.oval-forge.universal'
        else:
            self.namespace = namespace
        self.__db_init()
        self.__load_index(self.namespace)
    # Dropall
    def dropall(self):
        self.db.drop_collection('hashmap')
        self.db.drop_collection('index')
        self.db.drop_collection('oval_definition')
        self.db.drop_collection('oval_object')
        self.db.drop_collection('oval_variable')
        self.db.drop_collection('oval_state')
        self.db.drop_collection('oval_test')
    # Basic init
    def __db_init(self):
        # ID  -> Colection
        self.collections = {
            'def':'oval_definition',
            'obj':'oval_object',
            'ste':'oval_state',
            'tst':'oval_test',
            'var':'oval_variable'
        }
        # Names -> Classes
        self.ovalclass = {
            'oval_definition':oval_definition,
            'oval_object':oval_object,
            'oval_state':oval_state,
            'oval_test':oval_test,
            'oval_variable':oval_variable
        }
        # {hash:hash, id:id}
        self.hashmap = self.db['hashmap']
        self.__ensure_indexes()
    def __ensure_indexes(self):
        for collection in self.collections.values():
            self.db[collection].ensure_index('id', background=True)
        self.db['hashmap'].ensure_index('hash', background=True)
    # Basic operations
    def put(self, ovalobject):
        if isinstance(ovalobject, oval_object):
            if not ovalobject.namespace:
                if not isinstance(ovalobject, oval_definition):
                    ovalobject.namespace = self.namespace + "#" + self.namespace.split(".")[len(self.namespace.split("."))-1]
                else:
                    ovalobject.namespace = "http://oval.mitre.org/XMLSchema/oval-definitions-5"
            if not ovalobject.oval:
                ovalobject.oval = self.namespace
            ovalobject.clearme()
            object_hash = ovalobject.hashme()
            # If object in database
            in_base = self.hashmap.find_one({'class':ovalobject.__class__.__name__, 'hash':object_hash})
            if in_base:
                if hasattr(ovalobject, "datahash") and 'datahash' in in_base:
                    if ovalobject.datahash() == in_base['datahash']:
                        return in_base['id']
                    else:
                        in_base['datahash'] = ovalobject.datahash()
                        self.hashmap.update({'_id':in_base['_id']}, in_base, upsert=True)
                        ovalobject.id = in_base['id']
                        ovalobject.version = self.db[ovalobject.__class__.__name__].find_one({'id':in_base['id']})['version'] + 1
                        self.db[ovalobject.__class__.__name__].update({'id':in_base['id']}, ovalobject.serializeme())
                        return in_base['id']
                return in_base['id']
            # If it is not in DB
            else:
                if not ovalobject.id:
                    ovalobject.assign_id(self.index[ovalobject.oval][ovalobject.__class__.__name__])
                    self.index[ovalobject.oval][ovalobject.__class__.__name__] += 1
                self.db[ovalobject.__class__.__name__].update({'id':ovalobject.id},ovalobject.serializeme(), upsert=True)
                if hasattr(ovalobject, "datahash"):
                    data = {'id':ovalobject.id, 'class':ovalobject.__class__.__name__, 'hash':object_hash, 'datahash':ovalobject.datahash()}
                else:
                    data = {'id':ovalobject.id, 'class':ovalobject.__class__.__name__, 'hash':object_hash}
                self.hashmap.update({'id':ovalobject.id}, data, upsert=True)
                return ovalobject.id
    def get(self, id):
        if id and isinstance(id, basestring) and len(id.split(":")) >= 4:
            type = id.split(":")[2]
            if type in self.collections:
                in_base = self.db[self.collections[type]].find_one({'id':id})
                if in_base:
                    obj = self.__to_object(in_base, type)
                    return obj
    # TODO: Dont forget to save_all
    def close(self):
        self.__save_index(self.namespace)
    # Deserializer
    def __to_object(self, mongo_entry, type):
        object_class = self.ovalclass[self.collections[type]]()
        object_class.__init__()
        for attribute in mongo_entry.keys():
            if hasattr(object_class, attribute):
                if attribute == 'variables' or attribute == 'metadata' or attribute == 'criteria':
                    object_class.__setattr__(attribute, self.__deserialize(mongo_entry[attribute]))
                else:
                    object_class.__setattr__(attribute, mongo_entry[attribute])
        return object_class
    def __deserialize(self, varlist):
        if varlist and isinstance(varlist, list):
            output = list()
            for var in varlist:
                output.append(self.__deserialize_var(var))
            return output
        else:
            return list()
    def __deserialize_var(self, var):
        if var and isinstance(var, dict):
            ovalvar = variable()
            for key in var:
                if key == 'variables' and var[key]:
                    for var in var[key]:
                        ovalvar.variables.append(self.__deserialize_var(var))
                else:
                    ovalvar.__setattr__(key, var[key])
            return ovalvar
    # Default hashes
    def __init_default_hashes(self):
        self.default_hashmap = dict()
        new = oval_definition()
        self.default_hashmap[new.hashme()] = None
        new = oval_object()
        self.default_hashmap[new.hashme()] = None
        new = oval_state()
        self.default_hashmap[new.hashme()] = None
        new = oval_test()
        self.default_hashmap[new.hashme()] = None
        new = oval_variable()
        self.default_hashmap[new.hashme()] = None
    # Index operations
    def __load_index(self, namespace):
        #### Index format: {type:index_type, namespace:namespace, current:int}
        self.index = dict()
        self.index[namespace] = dict()
        index_collection = self.db['index']
        for index in index_collection.find({'namespace':namespace}):
            self.index[namespace][index['type']] = index['current']
        if not self.index[namespace]:
            self.index[namespace]['oval_definition'] = 0
            self.index[namespace]['oval_object'] = 0
            self.index[namespace]['oval_state'] = 0
            self.index[namespace]['oval_test'] = 0
            self.index[namespace]['oval_variable'] = 0
        return self.index
    def __save_index(self, namespace):
        if hasattr(self, 'index') and self.index and namespace in self.index:
            index_collection = self.db['index']
            for key in self.index[namespace]:
                if index_collection.find_one({'type':key, 'namespace':namespace}):
                    index = index_collection.find_one({'type':key, 'namespace':namespace})
                    index['current'] = self.index[namespace][key]
                    index_collection.update({'_id':index['_id']}, index, upsert=False)
                else:
                    index = {'type':key, 'namespace':namespace, 'current':self.index[namespace][key]}
                    index_collection.insert(index)
    # Global Functions from old lib
    # Object -> Etree
    def xml_me(self, object, easy = 0):
        if object and (isinstance(object, oval_object) or isinstance(object, variable)):
            if isinstance(object, variable):
                var = etree.Element(object.tag)
                if object.body:
                    var.text = object.body
                for k in object.attributes.keys():
                    var.attrib[k] = object.attributes[k]
                for sub in object.variables:
                    var.append(self.xml_me(sub))
                return var
            elif isinstance(object, oval_definition):
                definition = etree.Element(object.tag)
                definition.attrib['id'] = object.id
                if object.deprecated:
                    definition.attrib['deprecated'] = "true"
                definition.attrib['class'] = object.defclass
                definition.attrib['version'] = str(object.version)
                definition.attrib['xmlns'] = object.namespace
                metadata = etree.SubElement(definition,"metadata")
                etree.SubElement(metadata, "title").text = str(object.title)
                if object.comment:
                    definition.attrib['comment'] = object.comment
                if not easy:
                    etree.SubElement(metadata, "description").text = unicode(object.description)
                    for element in object.metadata:
                        metadata.append(self.xml_me(element))
                    for element in object.notes:
                        metadata.append(self.xml_me(element))
                if object.criteria:
                    for element in object.criteria:
                        definition.append(self.xml_me(element))
                return definition
            elif isinstance(object, oval_test):
                test = etree.Element(object.tag)
                test.attrib['id'] = object.id
                test.attrib['version'] = str(object.version)
                test.attrib['xmlns'] = object.namespace
                if object.deprecated:
                    test.attrib['deprecated'] = "true"
                if object.comment:
                    test.attrib['comment'] = object.comment
                if object.check:
                    test.attrib['check'] = object.check
                if object.check_existance:
                    test.attrib['check_existance'] = object.check_existance
                if object.state_operator:
                    test.attrib['state_operator'] = object.state_operator
                for element in object.variables:
                    test.append(self.xml_me(element))
                for element in object.notes:
                    test.append(self.xml_me(element))
                return test
            elif isinstance(object, oval_state):
                state = etree.Element(object.tag)
                state.attrib['id'] = object.id
                state.attrib['version'] = str(object.version)
                state.attrib['xmlns'] = object.namespace
                if object.operator:
                    state.attrib['operator'] = object.operator
                if object.deprecated:
                    state.attrib['deprecated'] = "true"
                if object.comment:
                    state.attrib['comment'] = object.comment
                for element in object.variables:
                    state.append(self.xml_me(element))
                for element in object.notes:
                    state.append(self.xml_me(element))
                return state
            elif isinstance(object, oval_variable):
                var = etree.Element(object.tag)
                var.attrib['id'] = object.id
                var.attrib['version'] = str(object.version)
                var.attrib['xmlns'] = object.namespace
                if object.datatype:
                    var.attrib['datatype'] = object.datatype
                if object.deprecated:
                    var.attrib['deprecated'] = "true"
                if object.comment:
                    var.attrib['comment'] = object.comment
                for element in object.variables:
                    var.append(self.xml_me(element))
                for element in object.notes:
                    var.append(self.xml_me(element))
                return var
            elif isinstance(object, oval_object):
                obj = etree.Element(object.tag)
                obj.attrib['id'] = object.id
                obj.attrib['version'] = str(object.version)
                obj.attrib['xmlns'] = object.namespace
                if object.deprecated:
                    obj.attrib['deprecated'] = "true"
                if object.comment:
                    obj.attrib['comment'] = object.comment
                for element in object.variables:
                    obj.append(self.xml_me(element))
                for element in object.notes:
                    obj.append(self.xml_me(element))
                return obj
        else:
            raise Exception("Only OVAL objects usable as input")
    # Indenter
    def indent(self, elem, level=0):
        i = "\n" + level*"  "
        if elem and len(elem):
            if not elem.text or not elem.text.strip():
                elem.text = i + "  "
            if not elem.tail or not elem.tail.strip():
                elem.tail = i
            for elem in elem:
                self.indent(elem, level+1)
            if not elem.tail or not elem.tail.strip():
                elem.tail = i
        else:

            if level and (not elem.tail or not elem.tail.strip()):
                elem.tail = i
    # Tag Splitter
    def SplitTagXmlns(self, xmlnstag):
        #return xmlns and simple tag name
        m = re.match(r"^\{(.*?)\}(.*)$",xmlnstag)
        if m is not None:
            return (m.group(1),m.group(2))
        else:
            return ("", xmlnstag)
    # Object -> String
    @lfu_cache(maxsize=100)
    def to_string(self, element):
        if (element and isinstance(element, oval_object)) or (element and isinstance(element, variable)):
            element = self.xml_me(element)
        if hasattr(element, 'tag') and hasattr(element, 'attrib'):
            if not element.text:
                element.text = ""
            for attribute in element.attrib:
                if not element.attrib[attribute]:
                    element.attrib[attribute] = ""
            new_tree = etree.ElementTree(element)
            self.indent(element)
            string = etree.tostring(element,"UTF-8").replace("<?xml version='1.0' encoding='UTF-8'?>", "").strip()
            del new_tree
            if string and isinstance(string,basestring):
                return string
        else:
            return None
    # Technical
    def _garbager(self, root):
        for element in list(root):
            self._garbager(element)
        root.clear()
        del root
    def _fssync(self, file):
        file.flush()
        os.fsync(file.fileno())
    # External XML parser
    def read_variables(self, etreeElement):
        new_var = variable(self.SplitTagXmlns(etreeElement.tag)[1], etreeElement.text, dict(etreeElement.attrib))
        for children in list(etreeElement):
            new_var.variables.append(self.read_variables(children))
        return new_var
    def read_object(self, etreeElement):
        if "id" in etreeElement.attrib.keys() and len(etreeElement.attrib['id'].split(":")) >= 4:
            # Find id and create new object according to it type
            id = etreeElement.attrib['id']
            type = etreeElement.attrib['id'].split(":")[2]
            namespace, tag = self.SplitTagXmlns(etreeElement.tag)
            if type == "def":
                oval_obj = oval_definition()
                oval_obj.oval = etreeElement.attrib['id'].split(":")[1]
                oval_obj.id = id
                oval_obj.tag = tag
                oval_obj.namespace = namespace
                if "class" in etreeElement.attrib.keys():
                    oval_obj.defclass = etreeElement.attrib['class']
                if "version" in etreeElement.attrib.keys() and str(etreeElement.attrib['version']).isdigit():
                    oval_obj.version = int(etreeElement.attrib['version'])
                if "comment" in etreeElement.attrib.keys():
                    oval_obj.comment = etreeElement.attrib['comment']
                if "deprecated" in etreeElement.attrib.keys():
                    oval_obj.deprecated = etreeElement.attrib['deprecated']
                # Build internals
                for children in list(etreeElement):
                    if self.SplitTagXmlns(children.tag)[1] == "metadata":
                        for metadata_el in list(children):
                            if self.SplitTagXmlns(metadata_el.tag)[1] == "title":
                                oval_obj.title = metadata_el.text
                            elif self.SplitTagXmlns(metadata_el.tag)[1] == "description":
                                oval_obj.description = metadata_el.text
                            elif self.SplitTagXmlns(metadata_el.tag)[1] == "note":
                                oval_obj.notes.append(self.read_variables(metadata_el))
                            else:
                                oval_obj.metadata.append(self.read_variables(metadata_el))
                    else:
                        oval_obj.criteria.append(self.read_variables(children))
                return oval_obj

            elif type == "obj":
                oval_obj = oval_object()
                oval_obj.oval = etreeElement.attrib['id'].split(":")[1]
                oval_obj.id = id
                oval_obj.tag = tag
                oval_obj.namespace = namespace
                if "version" in etreeElement.attrib.keys() and str(etreeElement.attrib['version']).isdigit():
                    oval_obj.version = int(etreeElement.attrib['version'])
                if "comment" in etreeElement.attrib.keys():
                    oval_obj.comment = etreeElement.attrib['comment']
                if "deprecated" in etreeElement.attrib.keys():
                    oval_obj.deprecated = etreeElement.attrib['deprecated']
                # Build internals
                for children in list(etreeElement):
                    if self.SplitTagXmlns(children.tag)[1] == "note":
                        oval_obj.notes.append(self.read_variables(children))
                    else:
                        oval_obj.variables.append(self.read_variables(children))
                return oval_obj

            elif type == "tst":
                oval_obj = oval_test()
                oval_obj.oval = etreeElement.attrib['id'].split(":")[1]
                oval_obj.id = id
                oval_obj.tag = tag
                oval_obj.namespace = namespace
                if "version" in etreeElement.attrib.keys() and str(etreeElement.attrib['version']).isdigit():
                    oval_obj.version = int(etreeElement.attrib['version'])
                if "comment" in etreeElement.attrib.keys():
                    oval_obj.comment = etreeElement.attrib['comment']
                if "deprecated" in etreeElement.attrib.keys():
                    oval_obj.deprecated = etreeElement.attrib['deprecated']
                if "check" in etreeElement.attrib.keys():
                    oval_obj.check = etreeElement.attrib['check']
                if "check_existance" in etreeElement.attrib.keys():
                    oval_obj.check_existance = etreeElement.attrib['check_existance']
                if "state_operator" in etreeElement.attrib.keys():
                    oval_obj.state_operator = etreeElement.attrib['state_operator']
                # Build internals
                for children in list(etreeElement):
                    if self.SplitTagXmlns(children.tag)[1] == "note":
                        oval_obj.notes.append(self.read_variables(children))
                    else:
                        oval_obj.variables.append(self.read_variables(children))
                return oval_obj

            elif type == "ste":
                oval_obj = oval_state()
                oval_obj.oval = etreeElement.attrib['id'].split(":")[1]
                oval_obj.id = id
                oval_obj.tag = tag
                oval_obj.namespace = namespace
                if "version" in etreeElement.attrib.keys() and str(etreeElement.attrib['version']).isdigit():
                    oval_obj.version = int(etreeElement.attrib['version'])
                if "comment" in etreeElement.attrib.keys():
                    oval_obj.comment = etreeElement.attrib['comment']
                if "deprecated" in etreeElement.attrib.keys():
                    oval_obj.deprecated = etreeElement.attrib['deprecated']
                if "operator" in etreeElement.attrib.keys():
                    oval_obj.operator = etreeElement.attrib['operator']
                # Build internals
                for children in list(etreeElement):
                    if self.SplitTagXmlns(children.tag)[1] == "note":
                        oval_obj.notes.append(self.read_variables(children))
                    else:
                        oval_obj.variables.append(self.read_variables(children))
                return oval_obj

            elif type == "var":
                oval_obj = oval_variable()
                oval_obj.oval = etreeElement.attrib['id'].split(":")[1]
                oval_obj.id = id
                oval_obj.tag = tag
                oval_obj.namespace = namespace
                if "version" in etreeElement.attrib.keys() and str(etreeElement.attrib['version']).isdigit():
                    oval_obj.version = int(etreeElement.attrib['version'])
                if "comment" in etreeElement.attrib.keys():
                    oval_obj.comment = etreeElement.attrib['comment']
                if "deprecated" in etreeElement.attrib.keys():
                    oval_obj.deprecated = etreeElement.attrib['deprecated']
                if "datatype" in etreeElement.attrib.keys():
                    oval_obj.datatype = etreeElement.attrib['datatype']
                # Build internals
                for children in list(etreeElement):
                    if self.SplitTagXmlns(children.tag)[1] == "note":
                        oval_obj.notes.append(self.read_variables(children))
                    else:
                        oval_obj.variables.append(self.read_variables(children))
                return oval_obj

            else:
                raise Exception("Error occured: Looks like not OVAL XML")
        else:
            raise Exception("Error occured: Wrong Etree element input")
    def parse_oval_xml(self, filename):
        if filename and isinstance(filename, basestring):
            xml_structured = oval_mapper(filename)

            for key in xml_structured.definitions.keys():
                self.put(self.read_object(xml_structured.definitions[key]))

            for key in xml_structured.objects.keys():
                self.put(self.read_object(xml_structured.objects[key]))

            for key in xml_structured.tests.keys():
                self.put(self.read_object(xml_structured.tests[key]))

            for key in xml_structured.states.keys():
                self.put(self.read_object(xml_structured.states[key]))

            for key in xml_structured.vars.keys():
                self.put(self.read_object(xml_structured.vars[key]))
        else:
            raise Exception("Must provide correct filename")
    # Dependencies searcher
    def get_variable_deps(self, var, pattern):
        if var and isinstance(var, variable) and pattern and isinstance(pattern, basestring):
            result = set()
            for value in var.attributes.values():
                if re.findall(pattern, value):
                    result = result.union(set(re.findall(pattern, value)))
            for child in var.variables:
                result = result.union(set(self.get_variable_deps(child, pattern)))
            return result
    @lfu_cache(maxsize=10000)
    def collect_dependencies(self, id, pattern = "(oval:.+?:.+?:\d+)"):
        if id and isinstance(id, basestring):
            ovalobject = self.get(id)
            if ovalobject:
                parsefield = 'variables'
                if hasattr(ovalobject, 'criteria'):
                    parsefield = 'criteria'
                result = set()
                for var in ovalobject.__getattribute__(parsefield):
                    result = result.union(self.get_variable_deps(var, pattern))
                for dep_id in result:
                    result = result.union(self.collect_dependencies(dep_id, pattern))
                return result
            else:
                return set()
        return set()
    # XML Exporter
    def produce_xml(self, filename, id_list = list(), ovalname = None, easy = 0):
        if isinstance(id_list, list) and isinstance(filename, basestring):
            output_file = open(filename, "w")
            output_file.write("<?xml version=\"1.0\" encoding=\"UTF-8\"?>" +"\n")
            output_file.write("<oval_definitions ")

            ns = dict()
            ns["xmlns"] = "http://oval.mitre.org/XMLSchema/oval-definitions-5"
            ns["xmlns:oval"] = "http://oval.mitre.org/XMLSchema/oval-common-5"
            ns["xmlns:oval-def"] = "http://oval.mitre.org/XMLSchema/oval-definitions-5"
            ns["xmlns:unix-def"] = "http://oval.mitre.org/XMLSchema/oval-definitions-5#unix"
            ns["xmlns:red-def"] = "http://oval.mitre.org/XMLSchema/oval-definitions-5#linux"
            ns["xmlns:xsi"] = "http://www.w3.org/2001/XMLSchema-instance"
            ns["xsi:schemaLocation" ] = "http://oval.mitre.org/XMLSchema/oval-common-5 oval-common-schema.xsd http://oval.mitre.org/XMLSchema/oval-definitions-5 oval-definitions-schema.xsd http://oval.mitre.org/XMLSchema/oval-definitions-5#unix unix-definitions-schema.xsd http://oval.mitre.org/XMLSchema/oval-definitions-5#linux linux-definitions-schema.xsd"

            if not ovalname:
                ovalname = self.namespace

            for key in ns.keys():
                output_file.write(key+"="+"\""+ns[key]+"\" ")
            output_file.write(">"+"\n")

            date = datetime.datetime.now()
            timestamp = str(date.strftime("%Y-%m-%dT%H:%M:%S"))+"\n"

            generator = etree.Element("generator")
            etree.SubElement(generator, "oval:product_name").text = "OVAL-FORGE Security Repository"
            etree.SubElement(generator, "oval:product_version").text = "1.0"
            etree.SubElement(generator, "oval:schema_version").text = "5.10"
            etree.SubElement(generator, "oval:timestamp").text = timestamp

            output_file.write(self.to_string(generator))
            self._garbager(generator)

            job = list()

            if not id_list:
                job = list()
                for collection in self.collections.values():
                    job_elements = self.db[collection].find({'oval':ovalname})
                    for n in job_elements:
                        job.append(n['id'])
            else:
                id_list = list(set(id_list))
                for id in id_list:
                    job.append(id)
                    job += self.collect_dependencies(id)

            job = list(set(job))

            definitions = list()
            tests = list()
            objects = list()
            states = list()
            vars = list()

            for key in job:
                if key.split(":")[2] == 'def':
                    definitions.append(key)
                elif key.split(":")[2] == 'ste':
                    states.append(key)
                elif key.split(":")[2] == 'tst':
                    tests.append(key)
                elif key.split(":")[2] == 'var':
                    vars.append(key)
                elif key.split(":")[2] == 'obj':
                    objects.append(key)

            output_file.write("<definitions>"+"\n")
            for key in definitions:
                xml = self.xml_me(self.get(key), easy)
                output_file.write(self.to_string(xml))
                self._garbager(xml)
                self._fssync(output_file)
            output_file.write("</definitions>"+"\n")

            output_file.write("<tests>"+"\n")
            for key in tests:
                xml = self.xml_me(self.get(key), easy)
                output_file.write(self.to_string(xml))
                self._garbager(xml)
                self._fssync(output_file)
            output_file.write("</tests>"+"\n")

            output_file.write("<objects>"+"\n")
            for key in objects:
                xml = self.xml_me(self.get(key), easy)
                output_file.write(self.to_string(xml))
                self._garbager(xml)
                self._fssync(output_file)
            output_file.write("</objects>"+"\n")

            output_file.write("<states>"+"\n")
            for key in states:
                xml = self.xml_me(self.get(key), easy)
                output_file.write(self.to_string(xml))
                self._garbager(xml)
                self._fssync(output_file)
            output_file.write("</states>"+"\n")

            output_file.write("<variables>"+"\n")
            for key in vars:
                xml = self.xml_me(self.get(key), easy)
                output_file.write(self.to_string(xml))
                self._garbager(xml)
                self._fssync(output_file)
            output_file.write("</variables>"+"\n")

            output_file.write("\n"+"</oval_definitions>")
            output_file.close()
        else:
            raise Exception("Filename for output not listed")
# Support classes for web operation
# Path splitter
os.chdir(os.path.dirname(sys.argv[0]))
# Main webkit class
class webkit():

    def __init__(self):
        # MITRE HOSTS
        self.mitre = dict()
        self.mitre['main'] = "mitre.org:80"
        self.mitre['cve'] = "cve.mitre.org:80"
        self.mitre['oval'] = "oval.mitre.org:443"
        # NVD HOSTS
        self.nvd = dict()
        self.nvd['main'] = "nvd.nist.gov:80"
        self.nvd['static'] = "static.nvd.nist.gov:80"

    @lfu_cache(maxsize=10)
    def get_mitre_oval(self, oval, force = 0):
        if not oval:
            return None
        if not os.path.exists("mitre_oval/"+ oval +".xml") or force:
            host = self.mitre['oval']
            link ="/rep-data/5.10/org.mitre.oval/v/platform/"+oval+".xml"
            self.save_url(host, link, "mitre_oval/"+ oval +".xml")
        oval_file = open("mitre_oval/"+oval+".xml", 'r')
        oval_data = oval_file.read()
        oval_file.close()
        if not oval_data and not force:
            return self.get_mitre_oval(oval, 1)
        return oval_data

    @lfu_cache(maxsize=100)
    def get_cve_mappings(self, os):
        if not os:
            return None
        os = os.upper()
        table_host = self.mitre['cve']
        table_link = "/data/refs/refmap/"
        table_index = "source-"+ os + ".html"
        table_src = self.get_url(table_host, table_link + table_index)
        if table_src:
            mappings = "\n".join(table_src.split("\n")[1:])
            mappings = re.findall("(?s)<tr>\n<td>"+ os + ":.*?</tr>", mappings)
            mappings = "<table>"+("\n".join(mappings))+"</table>"
            mappings = etree.fromstring(mappings)
            result = dict()
            for tablerow in mappings.findall("tr"):
                usn = ""
                cve = set()
                for td in tablerow.findall("td"):
                    if not list(td):
                        usn = td.text.replace(os + ":", "")
                    else:
                        for cvelink in td.findall("a"):
                            cve.add(cvelink.text)
                result[str(usn).lower()] = list(cve)
            return result

    @lfu_cache(maxsize=1000)
    def get_cve_data(self, cve, force = 0):
        # Checking or downloading neccesary CVE file
        if not cve or not isinstance(cve, basestring):
            return None
        cve = cve.split("-")
        if str(cve[1]).isdigit() and int(cve[1]) < 2000:
            return None
        if not os.path.exists("cve/nvdcve-2.0-"+cve[1]+".xml") or force:
            host = self.nvd['static']
            link ="/feeds/xml/cve/nvdcve-2.0-"+cve[1]+".xml"
            self.save_url(host, link, "cve/nvdcve-2.0-"+cve[1]+".xml")
        # Opening file
        datafile = open("cve/nvdcve-2.0-"+cve[1]+".xml", 'r')
        cvedata = datafile.read()
        datafile.close()
        # Finding needable CVE block
        data = dict()
        result = list()
        cve = "-".join(cve)
        cve_block = re.findall(r"(?s)<entry id=\""+cve+r"\">.+?</entry>", cvedata)
        del cvedata
        if not cve_block and not force:
            return_data = self.get_cve_data(cve, 1)
            return return_data
        elif not cve_block:
            return None
        # Cutting CPE list
        if cve_block:
            cve_block = cve_block[0]
            result = re.findall(r"(?s)<vuln:product>(.+?)</vuln:product>", cve_block)
        result = list(result)
        data['cpe_list'] = result
        # Cutting CVSS score and making vector
        result = dict()
        if cve_block:
            if len(re.findall(r"(?s)<cvss:score.*?>(.+?)</cvss:score.*?>", cve_block)):
                result['score'] = re.findall(r"(?s)<cvss:score.*?>(.+?)</cvss:score.*?>", cve_block)[0][0]
                result['vector'] = "AV:"+re.findall(r"(?s)<cvss:access-vector.*?>(.+?)</cvss:access-vector.*?>", cve_block)[0][0]+"/"
                result['vector'] += "AC:"+re.findall(r"(?s)<cvss:access-complexity.*?>(.+?)</cvss:access-complexity.*?>", cve_block)[0][0]+"/"
                result['vector'] += "Au:"+re.findall(r"(?s)<cvss:authentication.*?>(.+?)</cvss:authentication.*?>", cve_block)[0][0]+"/"
                result['vector'] += "C:"+re.findall(r"(?s)<cvss:confidentiality-impact.*?>(.+?)</cvss:confidentiality-impact.*?>", cve_block)[0][0]+"/"
                result['vector'] += "I:"+re.findall(r"(?s)<cvss:integrity-impact.*?>(.+?)</cvss:integrity-impact.*?>", cve_block)[0][0]+"/"
                result['vector'] += "A:"+re.findall(r"(?s)<cvss:availability-impact.*?>(.+?)</cvss:availability-impact.*?>", cve_block)[0][0]
            else:
                result['score'] = ""
                result['vector'] = ""
        else:
            result['score'] = ""
            result['vector'] = ""
        data['cvss'] = result
        # Cutting summary
        result = list()
        if cve_block:
            result = re.findall(r"(?s)<vuln:summary>(.+?)</vuln:summary>", cve_block)
        if len(result):
            result = result[0]
        else:
            result = None
        data['summary'] = result
        return data

    def get_url(self, host, link):
        host = host.split(":")
        connection_type = host[1]
        if connection_type == "443":
            url = "https://"+host[0]
        else:
            url = "http://"+host[0]
        url = url + link
        user_agent = 'Mozilla/4.0 (compatible; MSIE 5.5; Windows NT)'
        headers = { 'User-Agent' : user_agent }
        try:
            req = urllib2.Request(url, None, headers)
            response = urllib2.urlopen(req)
            htmlPage = response.read()
            if htmlPage and 'content-type' in response.headers.dict and response.headers.dict['content-type'] and response.headers.dict['content-type'].find("html") != -1:
                return htmlPage.decode("utf-8")
            else:
                return htmlPage
        except:
            return str()

    def save_url(self, host, link, filename):
        data = self.get_url(host, link)
        if data:
            if not os.path.exists(os.path.curdir  + "/" + "/".join(os.path.split(filename)[:len(os.path.split(filename))-1])):
                os.mkdir(os.path.curdir + "/" + "/".join(os.path.split(filename)[:len(os.path.split(filename))-1]))
            filename = os.path.join(os.path.dirname(os.path.curdir), filename)

            try:
                writedown = open(filename, 'wb')
                writedown.write(data)
                writedown.close()

            except:
                writedown = codecs.open(filename, 'wb', 'ISO-8859-15', 'replace')
                writedown.write(data)
                writedown.close()

            return filename

    def timestamp(self):
        date = datetime.datetime.now()
        return str(date.strftime("%Y-%m-%dT%H:%M:%S"))